"""来自论文Crater-based attitude and position estimation for planetary exploration with weighted measurement uncertainty
  不确定度的计算使用第3页
  加权POSIT算法的计算使用第7页

  TODO 复现出来有问题，感觉是作者的问题
"""

from .base import PoseEstimator, center_ellipse
import numpy as np


def lsq_ellipse_fit_uncertainty(x, y):
    x = x[:, np.newaxis]
    y = y[:, np.newaxis]
    A = np.hstack((2 * x * y, y * y, 2 * x, 2 * y, np.ones_like(x)))
    X = np.linalg.lstsq(A, -x * x, rcond=None)[0].squeeze()
    B, C, E, D, F = X
    # 公式8
    V = A @ X + (x * x).squeeze()
    # 不构造200*200的大矩阵，改用稀疏矩阵V*V实现公式10的R0
    CC = np.array([[1, X[0], X[2]], [X[0], X[1], X[3]], [X[2], X[3], X[4]]])
    x0, y0 = center_ellipse(CC)
    # 第4页公式12
    H = np.array(
        [
            [(E + 2 * B * x0), (-D - x0), -C, B, 0],
            [(D + 2 * B * y0), -y0, B, -1, 0],
        ]
    ) / (C - B * B)
    # 第4页公式10
    A = np.linalg.inv(A.T @ A) @ A.T
    P = V.var() * A @ A.T
    # R为误差不确定性矩阵
    R = H @ P @ H.T
    # 第4页公式13-15，即坐标的误差椭圆以及其倾角
    return R


class WP(PoseEstimator):
    def __init__(self, K, use_pnp_initial=True, *args, **kwargs):
        super().__init__(K, *args, **kwargs)
        self.use_pnp_initial = use_pnp_initial

    def __name__(self):
        return "wp"

    def forward(self, C_3d, C_2d, points_2d, *args, **ignore_kwargs):

        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        # 求解图像椭圆的误差矩阵
        Error_matrix = []
        for p_2d in points_2d:
            Error_matrix.append(lsq_ellipse_fit_uncertainty(p_2d[0], p_2d[1]))
        U, S, V = np.linalg.svd(np.linalg.inv(np.array(Error_matrix)))
        Q = np.einsum(
            "nij,njk->nik",
            np.sqrt(
                np.array([[S[:, 0], np.zeros(len(S))], [np.zeros(len(S)), S[:, 1]]])
            ).T,
            V,
        )
        # 调用opencv的库实现加权POSIT算法
        return self.posit_algorithm(X_cnt, x_cnt, Q)

    def posit_algorithm(
        self,
        points_3d,
        points_2d,
        Q,
        max_iterations=1000,
        tol=1e-8,
    ):
        """POSIT算法求解位姿"""
        R = np.eye(3)
        t = np.zeros(2)
        # 初始猜测
        lamda = np.ones(len(points_2d))
        # 迭代求解
        for _ in range(max_iterations):
            # 构建线性方程组
            A = np.block(
                [
                    [
                        Q[:, 0, 0, None] * points_3d,
                        Q[:, 0, 1, None] * points_3d,
                        Q[:, 0, 0, None],
                        Q[:, 0, 1, None],
                    ],
                    [
                        Q[:, 1, 0, None] * points_3d,
                        Q[:, 1, 1, None] * points_3d,
                        Q[:, 1, 0, None],
                        Q[:, 1, 1, None],
                    ],
                ]
            )
            b = np.block([lamda * points_2d[:, 0], lamda * points_2d[:, 1]])
            x = np.linalg.lstsq(A, b, rcond=None)[0]
            # s * c1, s * c2, s * t = x
            # 计算更新的比例因子
            s_now = np.sqrt(np.linalg.norm(x[:3]) * np.linalg.norm(x[3:6]))
            if np.linalg.norm(x[6:] / s_now - t) < tol:
                break
            # 计算当前迭代步的参数
            t = x[6:] / s_now
            Z_C = self.K[0, 0] / s_now
            R[0] = x[:3] / s_now
            R[1] = x[3:6] / s_now
            R[2] = np.cross(R[0], R[1])
            # 计算当前迭代数的lambda
            lamda = points_3d @ R[2] / Z_C + 1
        # 计算误差
        error = np.sqrt(np.power(A @ x - b, 2).mean())
        return True, R, np.block([t, Z_C])
